import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import glob

with np.load('F:\AI\data.jpg') as X:
    mtx,dist,rvecs,tvecs=[X[i] for i in ('mtx','dist','rvecs','tvecs')]

def draw(img,corners,imgpts):
    corner=tuple(corners[0].ravel())
    cv.line(img, corner, tuple(imgpts[0].ravel()),(255,0,0),5)
    cv.line(img, corner, tuple(imgpts[1].ravel()), (0, 255, 0), 5)
    cv.line(img, corner, tuple(imgpts[2].ravel()), (0, 0, 255), 5)
    return img

def draw2(img,imgpts):
    imgpts=np.int32(imgpts).reshape(-1,2)
    cv.drawContours(img,[imgpts[:4]],-1,(0,255,0),-3)
    for i,j in zip(range(4),range(4,8)):
        cv.line(img,tuple(imgpts[i]),tuple(imgpts[j]),255,3)
    cv.drawContours(img,[imgpts[4:]],-1,(0,0,255),3)
    return img
objp=np.zeros((8*7,3),np.float32)
objp[:,:2]=np.mgrid[0:7,0:8].T.reshape(-1,2)

axis=np.float32([[3,0,0],[0,3,0],[0,0,-3]]).reshape(-1,3)
images=glob.glob('F:\AI\data\left*.jpg')
for fname in images:
    img=cv.imread(fname)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    ret, corners = cv.findChessboardCorners(gray, (8, 7), None)
    if ret:
        rvecs,tvecs,inliers=cv.solvePnPRansac(objp,corners,mtx,dist)
        imgpts,jac=cv.projectPoints(axis,rvecs,tvecs,mtx,dist)
        img=draw(img,corners,imgpts)
        cv.imshow('img',img)
        k=cv.waitKey(0) & 0xff
        if k=='s':
            cv.imwrite('res.jpg',img)
print('ret:',ret)
print('内参数矩阵:\n',mtx)
print('畸变系数:\n',dist)
print('旋转向量:\n',rvecs)
print('平移向量:\n',tvecs)

cv.destroyAllWindows()